package com.yhy.common.security.token;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.NumberUtil;
import com.google.common.collect.Maps;
import com.yhy.common.redis.utils.RedisUtil;
import com.yhy.common.security.constant.SecurityConstants;
import com.yhy.common.security.model.MyUser;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;

import java.util.*;

/**
 * jwt token 扩展字段
 */
@Slf4j
public class JwtTokenEnhancer implements TokenEnhancer {

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        handler(accessToken, authentication);
        if (authentication.getPrincipal() instanceof MyUser) {
            MyUser user = (MyUser) authentication.getPrincipal();
            Map<String, Object> additionalInfo = Optional.ofNullable(accessToken.getAdditionalInformation()).orElse(new HashMap<>());
            additionalInfo.put(SecurityConstants.LICENSE, SecurityConstants.MY_LICENSE);
            additionalInfo.put(SecurityConstants.DETAILS_USER_ID, user.getId());
            additionalInfo.remove(SecurityConstants.DETAILS_USER);
            ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(additionalInfo);
        }
        return accessToken;
    }

    /**
     * jwt token有状态增强处理
     *
     * @param accessToken    accessToken
     * @param authentication authentication
     */
    private void handler(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        HashMap<String, Object> additionalInformation
                = Maps.newHashMap(Optional.ofNullable(accessToken.getAdditionalInformation()).orElse(new HashMap<>()));
        Object jti = additionalInformation.get(SecurityConstants.JWT_JTI);
        if (Objects.isNull(jti)) {
            return;
        }
        String key = SecurityConstants.CACHE_REDIS_ONLINE_USER + jti;
        Object expireSeconds = accessToken.getExpiration().toInstant().getEpochSecond();
        long seconds = NumberUtil.parseLong(expireSeconds.toString()) - DateUtil.currentSeconds();
        if (authentication.getPrincipal() instanceof MyUser) {
            MyUser user = ((MyUser) authentication.getPrincipal());
            String clientId = user.getClientId();
            Long userId = user.getId();
            additionalInformation.put(SecurityConstants.DETAILS_USER_ID, userId);
            additionalInformation.put(SecurityConstants.AUTHORIZATION_CLIENT, clientId);
        }
        String expirationFormat = DateUtil.formatDateTime(new Date(accessToken.getExpiration().toInstant().toEpochMilli()));
        additionalInformation.put(SecurityConstants.JWT_EXP, expirationFormat);
        additionalInformation.put(SecurityConstants.JWT_AUTHORITIES_KEY, authentication.getAuthorities());

        RedisUtil.hmSet(key, additionalInformation, seconds);
    }
}
